package com.lvhx.springboot.partitioner;

import org.springframework.batch.item.ExecutionContext;
import org.springframework.batch.item.ItemStreamException;
import org.springframework.batch.item.support.AbstractItemCountingItemStreamItemReader;
import org.springframework.util.ClassUtils;

import java.util.List;

public class CustomItemReader extends AbstractItemCountingItemStreamItemReader<Object> {
    private List<Object> data;
    private DataCenter<Object> dataCenter;
    private String dataKey;

    public CustomItemReader() {
        setName(ClassUtils.getShortName(CustomItemReader.class));
        setSaveState(false);
    }

    @Override
    protected Object doRead() throws Exception {
        if (data != null && data.size() > 0) {
            return (Object) data.remove(0);
        } else {
            return null;
        }
    }

    @Override
    protected void doOpen() throws Exception {

    }

    @Override
    protected void doClose() throws Exception {
        data = null;
		dataCenter.removeIfDataInMemory(dataKey);
        dataKey = null;
    }

    @Override
    public void open(ExecutionContext executionContext) throws ItemStreamException {
        dataKey = executionContext.getString("DATAKEY");

		data = dataCenter.getData(dataKey);

//        boolean immutableDataSource = (Boolean) executionContext
//                .get("IMMUTABLE");
//        if (immutableDataSource) {
//            setSaveState(true);
//        }

        super.open(executionContext);
    }

	public void setDataCenter(DataCenter<Object> dataCenter) {
		this.dataCenter = dataCenter;
	}
}
